# photo_add_hash.py
# python .\photo\photo_add_hash.py --db D:\baseline\photo --zip-file e:\

import os
import sqlite3
import hashlib
import zipfile
from pathlib import Path
from datetime import datetime
import argparse
import sys
from tqdm import tqdm

# 注册datetime适配器和转换器
def adapt_datetime(ts):
    return ts.isoformat()

def convert_datetime(s):
    return datetime.fromisoformat(s.decode('utf-8'))

sqlite3.register_adapter(datetime, adapt_datetime)
sqlite3.register_converter("DATETIME", convert_datetime)

def create_db(db_path):
    conn = sqlite3.connect(db_path, detect_types=sqlite3.PARSE_DECLTYPES)
    cursor = conn.cursor()
    
    # 创建表结构，注意create_time字段类型为DATETIME
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS photo_hash (
            id INTEGER PRIMARY KEY,
            zip_path TEXT NOT NULL,
            zip_name TEXT NOT NULL,
            file_path TEXT NOT NULL,
            file_name TEXT NOT NULL,
            file_hash TEXT NOT NULL,
            create_time DATETIME NOT NULL
        )
    ''')
    conn.commit()
    return conn

def calculate_file_hash(file_content):
    """计算给定文件内容的SHA256哈希值"""
    sha256_hash = hashlib.sha256()
    sha256_hash.update(file_content)
    return sha256_hash.hexdigest()

def scan_and_store_files(root_dir, zip_file):
    db_path = str(Path(root_dir) / 'file_hashes.db')
    conn = create_db(db_path)
    cursor = conn.cursor()
    
    try:
        print(f"正在处理ZIP文件: {zip_file}")
        with zipfile.ZipFile(zip_file, 'r') as archive:
            # 获取ZIP文件内的所有文件信息列表，并过滤掉目录项
            files_to_process = [file_info for file_info in archive.infolist() if not file_info.is_dir()]
    
            zip_file_path = Path(zip_file)
            relative_zip_path = str(zip_file_path.relative_to(root_dir).parent)  # 只获取ZIP文件的相对路径，不包括文件名
            zip_name = zip_file_path.name

            with tqdm(total=len(files_to_process), desc="处理文件") as pbar:
                for file_info in archive.infolist():
                    if not file_info.is_dir():  
                        with archive.open(file_info) as file:
                            file_content = file.read()
                            file_hash = calculate_file_hash(file_content)
                            # 提取文件名，不包括路径
                            just_filename = Path(file_info.filename).name
                            
                            # 检查哈希是否已存在于photo_hash表中
                            cursor.execute("SELECT COUNT(*) FROM photo_hash WHERE file_hash=?", (file_hash,))
                            if cursor.fetchone()[0] > 0:
                                print(f"文件哈希 {file_hash} 已存在，{just_filename} in {zip_name}")
                                continue
                            
                            current_time = datetime.now()  # 使用datetime对象
                            
                            cursor.execute('''
                                INSERT INTO photo_hash (zip_path, zip_name, file_path, file_name, file_hash, create_time)
                                VALUES (?, ?, ?, ?, ?, ?)
                            ''', (
                                relative_zip_path,  # 使用ZIP文件的相对路径（不包括文件名）
                                zip_name,
                                str(Path(file_info.filename).parent),
                                just_filename,  # 使用仅文件名
                                file_hash,
                                current_time  # 直接插入datetime对象
                            ))
                        pbar.update(1)
    except Exception as e:
        print(f"处理ZIP文件 {zip_file} 时出错: {e}")

    conn.commit()
    conn.close()

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Scan files under the specified directory and compute their hashes to be stored in an SQLite database")
    parser.add_argument('--db', required=True, help='Directory where the SQLite database is located')
    parser.add_argument('--zip-file', required=True, help='Zip file to scan')

    args = parser.parse_args()
    # 获取绝对路径，确保比较的是完整路径
    db_path_absolute = os.path.abspath(args.db)
    zip_file_absolute = os.path.abspath(args.zip_file)
    print(f"db:{db_path_absolute}; zip_file{zip_file_absolute}")
    # 检查 --db 目录是否存在
    if not os.path.isdir(db_path_absolute):
        print(f"Error: 目录 {db_path_absolute} 不存在")
        exit(1)

    # 检查 --zip-file 文件是否存在
    if not os.path.isfile(zip_file_absolute):
        print(f"Error: 文件 {zip_file_absolute} 不存在")
        exit(1)

    # 检查 --zip-file 是否是一个有效的ZIP文件
    try:
        with zipfile.ZipFile(zip_file_absolute, 'r') as zf:
            # 尝试读取ZIP文件的信息，如果文件不是有效的ZIP文件则会在此抛出异常
            zf.testzip()  # 此函数用于测试ZIP文件内的数据完整性，若返回None表示一切正常
    except zipfile.BadZipFile:
        print(f"Error: 文件 {zip_file_absolute} 不是有效的ZIP文件")
        exit(1)

    # 检查 --zip-file 的目录部分是否以 --db 开头
    zip_file_dir_absolute = os.path.dirname(zip_file_absolute)
    if not zip_file_dir_absolute.startswith(db_path_absolute):
        print("Error: --zip-file 所在的目录必须以 --db 参数的值开头")
        exit(1)

    db_path = os.path.join(args.db, 'file_hashes.db')
    zip_file = args.zip_file

    
    scan_and_store_files(args.db, zip_file)